from openai import OpenAI
import os
# 加载环境变量
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

client = OpenAI()


# 封装 OpenAI 的 Embedding 模型接口
#
# 预习 Q : Text Embedding: 这部分指代“文本嵌入”，这是一种在自然语言处理（NLP）中广泛使用的技术，
# 用于将文本数据（如词语、短语、句子或文档）转换为连续的、低维的向量表示。
# 这些向量保留了原始文本的语义信息，
# 使得计算机能够理解和处理文本数据，进而应用于诸如文本分类、情感分析、语义相似度计算、机器翻译等各种下游任务。
def get_embeddings(texts, model="text-embedding-ada-002", dimensions=None):
    # text-embedding-ada-002 这个模型不支持自定义维度?
    if model == "text-embedding-ada-002":
        dimensions = None
    if dimensions:
        data = client.embeddings.create(
            input=texts, model=model, dimensions=dimensions).data
    else:
        data = client.embeddings.create(
            input=texts, model=model).data
    # print("调用模型借口获取的 data{}".format(data))
    return [x.embedding for x in data]


# 封装 openai 接口
def get_completion(prompt, model="gpt-3.5-turbo-1106"):
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content
